(*
 *  MTASC - MotionTwin ActionScript2 Compiler
 *  Copyright (c)2004 Nicolas Cannasse
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

open Expr
open Swf
open ExtHashtbl
open ExtString
open ExtList
open GenSwfJump
open GenSwfTools
open GenSwfGenerator
open GenSwfGeneratorTools
open Printf

let enable_main = ref false
let version = ref None
let ftrace = ref None

let to_utf8 str =
	try
		UTF8.validate str;
		str;
	with
		UTF8.Malformed_code -> 
			let b = UTF8.Buf.create 0 in
			String.iter (fun c -> UTF8.Buf.add_char b (UChar.of_char c)) str;
			UTF8.Buf.contents b

let use_components = ref false
let separate = ref true
let keep = ref false
let bgcolor = ref 0xFFFFFF
let frame = ref 1
let header = ref None
let excludes = Hashtbl.create 0

let is_excluded (path,name) =
	if Hashtbl.mem excludes (s_type_path (path,name)) then
		true
	else
		let rec loop = function
			| [] -> false
			| p :: l ->
				Hashtbl.mem excludes (String.concat "." (List.rev ("*" :: p :: l))) || loop l
		in
		loop (List.rev path)

let support_overlength_classes = ref true

let generate file out ~compress exprs =
	let file , linkage =
		(try
			let f,l = String.split file "@" in
			f , Some l
		with
			Invalid_string ->
				file , None)
	in
	let ctx = {
		main = ref None;
		idents = Hashtbl.create 0;
		ops = DynArray.create();
		super_bindings = Hashtbl.create 0;
		current = Class.empty;
		code_pos = 1;
		ident_count = 0;
		stack = 0;
		reg_count = 0;
		locals = Hashtbl.create 0;
		stack_size = 0;
		cur_block = (EBreak,null_pos);
		breaks = [];
		forins = 0;
		continue_pos = 0;
		opt_push = false;
		curmethod = "";
	} in
	DynArray.add ctx.ops (AStringPool []);
	
	let tags = ref [] in
	let hpackages = Hashtbl.create 0 in
	Class.generate (fun clctx ->
		ctx.current <- clctx;
		let ctx = (if !separate then new_context ctx else ctx) in
		if not (Class.intrinsic clctx) && not (is_excluded (Class.path clctx)) then begin
			if !separate then DynArray.add ctx.ops (AStringPool []);
			let ssize = ActionScript.actions_length ctx.ops in
			generate_class_code ftrace enable_main ctx clctx (if !separate then Hashtbl.create 0 else hpackages);
			if !separate then tags := ("__Packages." ^ s_type_path (Class.path clctx),ctx.idents,ctx.ops) :: !tags;
			let size = ActionScript.actions_length ctx.ops in
			if size - ssize >= 1 lsl 15 then
			begin
				if !support_overlength_classes then
					prerr_endline ("Warning: Class " ^ s_type_path (Class.path clctx) ^ " has 32 KiB or more of bytecode, which is not supported by older MTASC versions.")
				else
					failwith ("Error: Class " ^ s_type_path (Class.path clctx) ^ " has 32 KiB or more of bytecode; please refactor it or recompile without -32k-limit.");
				eliminate_overlength_jumps ctx.ops
				(* Note: It is possible here for the UInt16 parameters in (for example) function/try/catch/finally/with instructions to be overlength too.
				   However, we don't support dealing with that, since it should be way less common in real-world code than even overlength
				   classes, which are already very uncommon. For that same reason, we also don't bother to output nice error messages
				   in such a case; the user will just see an Overflow exception from IO.write_ui16. If we ever do want to check for overlength
				   UInt16's, though, this would be the spot to do it. *)
			end;
		end;	
	) exprs;
(*
	if not (Sys.file_exists out) then print_endline "File not exists. after Class.generate";
*)
	if not !separate then tags := ("__Packages.MTASC",ctx.idents,ctx.ops) :: !tags;
		
	(match !(ctx.main) with
	| None ->
		if !enable_main then failwith "Main entry point not found";
	| Some (p,clname) -> 
		let ctx = new_context ctx in 
		DynArray.add ctx.ops (AStringPool []);
		(*//    (main class).main(this); *)
		push ctx [VStr "MTASC_MAIN"];
		write ctx (ASetReg 0);
		write ctx APop;
		push ctx [VStr "this"];
		write ctx AEval;
		push ctx [VInt 1];
		let k = generate_package ~fast:true ctx p in
		push ctx [VStr clname];
		getvar ctx k;
		push ctx [VStr "main"];
		call ctx VarObj 1;
		write ctx APop;
		tags := ("",ctx.idents,ctx.ops) :: !tags;
	);
	
	tags := List.rev !tags;
	List.iter (fun (n,idents,ops) ->
		let idents = Hashtbl.fold (fun ident pos acc -> (ident,pos) :: acc) idents [] in
		let idents = List.sort ~cmp:(fun (_,p1) (_,p2) -> compare p1 p2) idents in
		DynArray.set ops 0 (AStringPool (List.map (fun (id,_) -> to_utf8 id) idents));
	) !tags;
	let tag ?(ext=false) d = {
		tid = 0;
		textended = ext;
		tdata = d;
	} in
	let header , data = (match !header with
		| None ->
			let ch = IO.input_channel (open_in_bin file) in
			let header, data = (try Swf.parse ch with IO.No_more_input | IO.Overflow _ | IO.Bits_error -> failwith "Input swf is corrupted") in
			IO.close_in ch;
			header , data
		| Some h ->
			let data = [tag (TSetBgColor { cr = !bgcolor lsr 16; cg = (!bgcolor lsr 8) land 0xFF; cb = !bgcolor land 0xFF }) ] in
			let data = data @ (Array.to_list (Array.init !frame (fun _ -> tag TShowFrame))) in
			h , data)
	in
	let header = (match !version with None -> header | Some v -> { header with h_version = v }) in
	let found = ref false in
	let curf = ref !frame in
	let regs = ref [] in
	let found_ids = ref [] in
	let insert loop showf acc l =
		if !found || !curf > 1 then begin
			curf := !curf - 1;
			loop (showf @ acc) l
		end else begin
			found := true;
			let main = ref None in
			let rec loop_tag cid l = 
				if List.exists ((=) cid) !found_ids then
					loop_tag (cid + 1) l
				else
					loop_tag_rec cid l
			and loop_tag_rec cid = function				
				| [] -> []
				| ("",_,ops) :: l ->
					main := Some (tag ~ext:true (TDoAction ops));
					loop_tag (cid + 1) l
				| (name,_,ops) :: l ->
					tag ~ext:true (TClip { c_id = cid; c_frame_count = 1; c_tags = [] }) ::
					tag ~ext:true (TExport [{ exp_id = cid; exp_name = name }]) ::
					tag ~ext:true (TDoInitAction { dia_id = cid; dia_actions = ops }) ::
					loop_tag (cid + 1) l
			in
			let t = List.rev (loop_tag 0x5000 !tags) in
			loop (showf @ (match !main with None -> [] | Some m -> [m]) @ !regs @ t @ acc) l
		end
	in	
	let replace_package p cid x y z = 
		if p = "__Packages.MTASC" || (not !use_components && not !keep) then
			[]
		else try 			
			let t = List.find (fun (n,_,_) -> p = n) !tags in
			tags := List.filter ((!=) t) !tags;
			[x;y;tag ~ext:true (TDoInitAction { dia_id = cid; dia_actions = (match t with (_,_,o) -> o) })]
		with Not_found ->
			if !use_components && String.length p > 14 && String.sub p 0 14 = "__Packages.mx." then
				[x;y;z]
			else if !keep then
				[x;y;z]
			else
				[]
	in
	let rec loop acc = function
		| [] ->
			if not !found then failwith ("Frame " ^ string_of_int !frame ^ " not found in SWF");	
			List.rev acc
		| ({ tdata = TDoAction a } as x1) :: ({ tdata = TShowFrame } as x2) :: l ->
			if DynArray.length a > 0 && (match DynArray.get a 0 with AStringPool ("MTASC_MAIN" :: _) -> true | _ -> false) then
				loop acc (x2 :: l)
			else
				insert loop [x2;x1] acc l
		| ({ tdata = TShowFrame } as x) :: l ->
			insert loop [x] acc l
		| ({ tdata = TClip _ } as x) :: ({ tdata = TExport [{ exp_name = e; exp_id = cid }] } as y) :: ({ tdata = TDoInitAction _ } as z) :: l ->
			let l2 = replace_package e cid x y z in
			if l2 <> [] then found_ids := cid :: !found_ids;
			loop ((List.rev l2) @ acc) l
		| { tdata = TDoInitAction { dia_actions = d } } as x :: l ->
			let process mcname clname =
				let cpath = (match List.rev (String.nsplit clname ".") with [] -> assert false | x :: l -> List.rev l , x) in
				print_endline("during loop acc1");
				(try
					ignore(Class.getclass ctx.current cpath)
				with
				_ -> 
					if not !use_components || (match cpath with ("mx" :: _, _) -> false | _ -> true) then
						prerr_endline ("Warning : The MovieClip " ^ mcname ^ " needs the class " ^ clname ^ " which was not compiled :\nPlease force compilation of this class by adding it to the commandline."));				
				if !found then
					loop (x :: acc) l
				else begin
					regs := x :: !regs;
					loop acc l
				end
			in
			(match DynArray.to_list d with
			| [
				APush [PString clname];
				AEval;
				APush [PString mcname;PInt _;PString "Object"];
				AEval;
				APush [PString "registerClass"];
				AObjCall;
				APop
			] ->
				process mcname clname
			| [
				AStringPool [clname;"Object";"registerClass"];
				APush [PStack 0];
				AEval;
				APush [PStack 0;PInt _;PStack 1];
				AEval;
				APush [PStack 2];
				AObjCall;
				APop;
			] -> 				
				process clname clname
			| _ ->
				loop (x :: acc) l);
		| x :: l ->
			loop (x :: acc) l
	in
	let ch = IO.output_channel (open_out_bin out) in
	let t = loop [] data in
	Swf.write ch (header,t);
	IO.close_out ch

let make_header s =
	let sl = String.nsplit s ":" in
	try
		let make w h fps =
			let w = int_of_string w in
			let h = int_of_string h in
			{
				h_version = 7;
				h_size = {
					rect_nbits = if (max w h) >= 820 then 16 else 15;
					left = 0;
					top = 0;
					right = w * 20;
					bottom = h * 20;
				};
				h_frame_count = 1;
				h_fps = to_float16 (float_of_string fps);
				h_compressed = true;
			}
		in
		match sl with
		| [w;h;fps] ->
			make w h fps
		| [w;h;fps;color] ->
			bgcolor := int_of_string ("0x" ^ color);
			make w h fps;
		| _ ->
			raise Exit
	with
		_ -> raise (Arg.Bad "Invalid header format")

let rec trim f =
	let l = String.length f in
	if l = 0 then
		""
	else match f.[l - 1] with
		| '\r' | '\n' -> trim (String.sub f 0 (l - 1))
		| _ -> f

let exclude_file f =
	let lines = (try
		let ch = open_in (Plugin.find_file f) in
		let lines = Std.input_list ch in
		close_in ch;
		lines
	with Not_found | Sys_error _ ->
		String.nsplit f ";"
	) in
	List.iter (fun f ->
		let f = trim f in
		if f <> "" then Hashtbl.replace excludes f ()
	) lines

;;
generate_function_ref := generate_function ftrace;
SwfParser.init SwfZip.inflate SwfZip.deflate;
SwfParser.full_parsing := false; (* faster, safer *)
Swf.warnings := false;
let swf = ref None in
let out = ref None in
Plugin.add [
	("-swf",Arg.String (fun f -> swf := Some f),"<file> : swf file to update");
	("-out",Arg.String (fun f -> out := Some f),"<file> : swf output file");
	("-keep",Arg.Unit (fun () -> keep := true),": does not remove AS2 classes from input SWF");
	("-frame",Arg.Int (fun i -> if i <= 0 then raise (Arg.Bad "Invalid frame"); frame := i),"<frame> : export into target frame (must exist in the swf)");
	("-main",Arg.Unit (fun () -> enable_main := true),": enable main entry point");
	("-header",Arg.String (fun s -> header := Some (make_header s)),"<header> : specify header format 'width:height:fps'");
	("-group",Arg.Unit (fun () -> separate := false),": group classes into a single clip");
	("-exclude",Arg.String (fun f -> exclude_file f),"<file> : exclude classes listed in file");
	("-version",Arg.Int (fun n -> version := Some n),": change SWF version (6,7,8,...)");	
	("-trace",Arg.String (fun t -> ftrace := Some t),"<function> : specify a TRACE function");
	("-32k-limit",Arg.Unit (fun () -> support_overlength_classes := false),": treat overlength classes (>= 32 KiB of bytecode) as an error");
]
(fun t ->
	if !keep && !header <> None then failwith "-keep cannot be used together with -header";
	if !Plugin.verbose && Hashtbl.length excludes > 0 then Printf.printf "Excludes : %s\n" (String.concat ";" (List.of_enum (Hashtbl.keys excludes)));
	match !swf with 
	| None -> () 
	| Some f -> generate f (match !out with None -> f | Some f -> f) ~compress:true (Typer.exprs t) 
);
